import json
import numpy as np
import pandas as pd
import ast
from tqdm import tqdm

def operate_data():
    items = set()
    for mode in ['train', 'test']:
        print(f"--------{mode}数据处理--------")
        names=['id','uid', 'item_seq', 'pos_item', 'neg_item', 'pos_user_seq', 'neg_user_seq']
        data = pd.read_csv(f"{mode}_corpus_total_dual.txt", sep='\t',  names=names,
                           usecols=['uid',  'item_seq','pos_item','neg_item'])

        data['item_seq'] = data['item_seq'].apply(ast.literal_eval)
        for i in data.loc[:, 'item_seq']:
            items.update(i)

        pos_data = data.loc[:, ['uid', 'item_seq' ,'pos_item']]
        pos_data['behavior'] = 1
        pos_data.rename(columns={"pos_item":"target_item"}, inplace=True)

        neg_data = data.loc[:, ['uid', 'item_seq' ,'neg_item']]
        neg_data['behavior'] = 0
        neg_data.rename(columns={"neg_item":"target_item"}, inplace=True)

        final_data = pd.concat([pos_data, neg_data], axis=0).sort_values('uid').reset_index(drop=True)

        final_data.to_csv(f"final_{mode}_data.csv", sep='\t', index=False)

    print("--------映射文件存储--------")
    items = sorted(list(items))

    id2itemid = {
        int(id) +1:int(itemid) for id, itemid in enumerate(items)
    }
    itemid2id = {
        int(itemid):int(id)+1 for id, itemid in enumerate(items)
    }

    with open('itemid2idx.json', 'w', encoding='utf-8') as f:
        json.dump(itemid2id, f, ensure_ascii=False, indent=4)
    with open('idx2itemid.json', 'w', encoding='utf-8') as f:
        json.dump(id2itemid, f, ensure_ascii=False, indent=4)

def data_info():
    data = pd.read_csv("final_train_data.csv",sep='\t')
    data['item_seq'] = data['item_seq'].apply(ast.literal_eval)
    lens = []
    for l in tqdm(data.loc[:, 'item_seq']):
        lens.append(len(l))
    print(f"行为序列平均长度:{np.mean(lens)}")

if __name__ == '__main__':
    operate_data()
    data_info()


